/*******************************************************************************
 * CopyRight (c) 2020 
 *
 * \file    bsf.c
 * \brief   bit stream filter parser
 * \author  ZhangJing
 * \version 1.0.0
 * \date    2020-07-08
 
 * -----------------------------------------------------------------------------
 * Modify history：
 * <Date>       | <Version> | <Author>  | <Discribe>
 * 2020-07-08   | v1.0.0    | ZhangJing | Create file
 * -----------------------------------------------------------------------------
 ******************************************************************************/

#include <stdio.h>
#include <string.h>
#include <stdlib.h>

#include "publicdef.h"
#include "print_ctrl.h"
#include "format.h"

enum {
    H264_NAL_SLICE           = 1,
    H264_NAL_IDR_SLICE       = 5,
    H264_NAL_SPS             = 7,
    H264_NAL_PPS             = 8,
};

#define ADTS_HEADER_SIZE 7

/*!
 * from mp4 mdat to h264 stream
 */
int h264_mp4toannexb(void *h, char *ibuf, int ilen, char *obuf, int *olen)
{
    FormatContext *fc = h;
    Stream *st = fc->streams[0];
    char *input_buf = ibuf;
    unsigned int tmp;
    unsigned char unit_type, sps_seen = 0, new_idr, pps_seen = 0;
    unsigned long out_size = 0;
    unsigned int start_code4 = 0x01000000;
    char start_code3[3] = {0x00, 0x00, 0x01};
    char *codec_private = NULL;
    int private_len = 0;

    //NOTE: got extradata form segment/tracks/codecprivate(0x63a2)
    if (st->mkvst) {
        if (st->mkvst->pps_data)
            codec_private = st->mkvst->pps_data;
        private_len = st->mkvst->pps_data_len;
    } else { /* mp4 */
        if (st->mp4st->pps_data)
            codec_private = st->mp4st->pps_data;
        private_len = st->mp4st->pps_data_len;
    }

    //FIXME: new_idr = H264BSFContext->idr_sps_seen
    new_idr = 1;

    *olen = 0;
    do {
        unsigned int nal_size = 0;
        /* get length */
        memcpy(&tmp, ibuf, sizeof(tmp));
        ibuf += sizeof(tmp);
        nal_size = MKLE32(tmp);
        pr_detail("h264 nal_size:%d", nal_size);
        if (nal_size > ilen) {
            pr_err("nal_size: %d(0x%08x) error!, tmp: 0x%08x",
                nal_size, nal_size, tmp);
            *olen = 0;
            return -1;
        }
        if (!nal_size)
            continue;

        /* get h264 NAL type */
        unit_type = *ibuf & 0x1f;
        if (unit_type == H264_NAL_SPS) {
            sps_seen = new_idr = 1;
        } else if (unit_type == H264_NAL_PPS) {
            pps_seen = new_idr = 1;
            /* if SPS has not been seen yet, prepend the AVCC one to PPS */
            if (!sps_seen) {
                //TODO: write sps, sps_size
                sps_seen = 1;
            }
        }

        /* If this is a new IDR picture following an IDR picture, reset the idr flag.
            * Just check first_mb_in_slice to be 0 as this is the simplest solution.
            * This could be checking idr_pic_id instead, but would complexify the parsing. */
        if (!new_idr && unit_type == H264_NAL_IDR_SLICE && (ibuf[1] & 0x80))
            new_idr = 1;

        /* prepend only to the first type 5 NAL unit of an IDR picture, if no sps/pps are already present */
        if (new_idr && unit_type == H264_NAL_IDR_SLICE && !sps_seen && !pps_seen) {
            memcpy(obuf, codec_private, private_len);
            obuf += private_len;
            *olen += private_len;
            new_idr = 0;
        } else if (new_idr && unit_type == H264_NAL_IDR_SLICE && sps_seen && !pps_seen) {
            //TODO: write sps, sps_size
        }

        if (out_size && !(unit_type == H264_NAL_SPS || unit_type == H264_NAL_PPS)) {
            memcpy(obuf, start_code3, sizeof(start_code3));
            obuf += sizeof(start_code3);
            *olen += sizeof(start_code3);
            memcpy(obuf, ibuf, nal_size);
            obuf += nal_size;
            *olen += nal_size;
            ibuf += nal_size;
        }
        if ((unit_type == H264_NAL_SPS || unit_type == H264_NAL_PPS)
            || (!(unit_type == H264_NAL_SPS || unit_type == H264_NAL_PPS)
                && !out_size)) {
            memcpy(obuf, &start_code4, sizeof(start_code4));
            obuf += sizeof(start_code4);
            *olen += sizeof(start_code4);
            memcpy(obuf, ibuf, nal_size);
            obuf += nal_size;
            *olen += nal_size;
            ibuf += nal_size;
        }

        if (!new_idr && unit_type == H264_NAL_SLICE) {
            new_idr  = 1;
            sps_seen = 0;
            pps_seen = 0;
        }
    } while (ibuf < (input_buf + ilen));

    return 0;
}

/*!
 * from mp4 mdat to hevc/h265 stream
 */
static int hevc_mp4toannexb(void *h, char *ibuf, int ilen, char *obuf, int *olen)
{
    FormatContext *fc = h;
    Stream *st = fc->streams[0];
    unsigned int tmp;
    int got_irap = 0;
    unsigned int start_code4 = 0x01000000;
    //NOTE: got extradata from stsd box of mp4
    char *extradata = NULL;
    int extradata_len = 0;

    if (st->mp4st) {
        if (st->mp4st->pps_data)
            extradata = st->mp4st->pps_data;
        extradata_len = st->mp4st->pps_data_len;
    }

    *olen = 0;

    while (ilen) {
        unsigned int nalu_size = 0;
        int      nalu_type;
        int is_irap, add_extradata, extra_size, prev_size;

        memcpy(&tmp, ibuf, sizeof(tmp));
        nalu_size = MKLE32(tmp);
        ilen -= sizeof(tmp);
        ibuf += sizeof(tmp);
        nalu_type = (ibuf[0] >> 1) & 0x3f;

        /* prepend extradata to IRAP frames */
        is_irap       = nalu_type >= 16 && nalu_type <= 23;
        add_extradata = is_irap && !got_irap;
        extra_size    = add_extradata * sizeof(extradata);
        got_irap     |= is_irap;


        if (extra_size) {
            memcpy(obuf, extradata, extradata_len);
            obuf += extradata_len;
            *olen += extradata_len;
        }

        memcpy(obuf, &start_code4, sizeof(start_code4));
        obuf += sizeof(start_code4);
        *olen += sizeof(start_code4);

        memcpy(obuf, ibuf, nalu_size);
        obuf += nalu_size;
        *olen += nalu_size;
        ibuf += nalu_size;
        ilen -= nalu_size;
    }

    return 0;
}

/*!
 * \brief put adts to aac stream
 */
static int adts_aac(void *h, char *ibuf, int ilen, char *obuf, int *olen)
{
    unsigned int full_frame_size = ADTS_HEADER_SIZE + ilen;
    //FIXME: get objecttype, sample_rate_index and channel_conf value
    unsigned char channel_conf = 0x2;       //010
    unsigned char objecttype = 0x1;         //01
    unsigned char sample_rate_index = 0x4;  //0100

    /* 12 bits 0xfff, syncword*/
    obuf[0] = 0xff;
    obuf[1] = 0xf << 4;

    /* 1 bit 0, ID */
    /* 2 bits 0, layer */
    /* 1 bit 1, protection_absent */
    obuf[1] |= 0x01;

    /* 2 bits 01, objecttype */ //FIXME:
    /* 4 bits sample_rate_index, 0100 */ //FIXME:
    /* 1 bit 0 private_bit */
    /* 1 of 3 bits 010, channel_conf */ //FIXME:
    obuf[2] = (objecttype << 6) | (sample_rate_index << 2) | (channel_conf >> 2);

    /* 2 0f 3 bits 010, channel_conf */ //FIXME:
    /* 1 bit 0, original_copy */
    /* 1 bit 0, home */
    /* 1 bit 0, copyright_identification_bit */
    /* 1 bit 0, copyright_identification_start */
    /* 2 of 13 bits xxxx, full_frame_size */
    obuf[3] = (channel_conf << 6) | (full_frame_size >> (13 - 2));

    /* 8 0f 13 bits xxxx, full_frame_size */
    obuf[4] = ((full_frame_size >> 3) & 0xFF);

    /* 3 0f 13 bits xxxx, full_frame_size */
    /* 5 of 11 bits 0x7ff, adts_buffer_fullness */
    obuf[5] = (full_frame_size & 0x7) << 5;
    obuf[5] |= (0x7FF >> (11 - 5));

    /* 6 of 11 bits 0x7ff, adts_buffer_fullness */
    /* 2 bits 0, number_of_raw_data_blocks_in_frame */
    obuf[6] = (0x7ff >> (11 - 8)) & 0xFC;

    *olen += ADTS_HEADER_SIZE;
    memcpy(obuf + ADTS_HEADER_SIZE, ibuf, ilen);
    *olen += ilen;

    return 0;
}

BitStreamFilter h264_mp4toannexb_bsf = {
    "h264_bsf", ID_H264,
    h264_mp4toannexb
};

BitStreamFilter hevc_mp4toannexb_bsf = {
    "hevc_bsf", ID_HEVC,
    hevc_mp4toannexb
};
BitStreamFilter adts_bsf = {
    "adts_bsf", ID_AAC,
    adts_aac
};
